#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import random
import math
import copy
import matplotlib.pyplot as plt
import pylab as mpl

mpl.rcParams['font.sans-serif'] = ['SimHei']

########################################################################
#          1) 从txt中读取“初始参数”和“设备信息” + 全局变量准备           #
########################################################################

def read_initial_params(params_file):
    """
    读取“初始参数.txt”中：
    1) 第1行是注释，不管
    2) 第2行: 迭代次数
    3) 第3行: 种群个数
    4) 第4行: 总长 L
    5) 第5行: 总宽 W
    6) 第6行: 总高 H
    7) 第7行: 大层数 C
    8) 第8行: 两边小层数 c1
    9) 第9行: 中间小层数 c2
    10) 第10行: 两边长度 a
    11) 第11行: 中间长度 b
    12) 第12行: 设备最小间距 d_min
    """
    with open(params_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    # 去掉空行、换行
    lines = [l.strip() for l in lines if l.strip()]

    # 逐行解析
    max_iter = int(lines[1])
    pop_size = int(lines[2])
    L = float(lines[3])
    W = float(lines[4])
    H = float(lines[5])
    C = int(lines[6])
    c1 = int(lines[7])
    c2 = int(lines[8])
    a = float(lines[9])
    b = float(lines[10])
    d_min = float(lines[11])

    return max_iter, pop_size, L, W, H, C, c1, c2, a, b, d_min


def read_devices_data(data_file):
    """
    读取“数据.txt”中 23 行设备信息 (第1行为表头)。
    每行: 设备ID  长  宽  高  重量  朝向v  x  y  z
    返回 list of dict
    """
    with open(data_file, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    lines = [l.strip() for l in lines if l.strip()]

    # 第1行表头, 后面是数据
    data_lines = lines[1:]
    devices_info = []

    for line in data_lines:
        parts = line.split()
        dev_id = int(parts[0])
        L_ = float(parts[1])
        W_ = float(parts[2])
        H_ = float(parts[3])
        mass = float(parts[4])
        v_ = int(parts[5])
        x_ = float(parts[6])
        y_ = float(parts[7])
        z_ = float(parts[8])

        dev_dict = {
            'id': dev_id,
            'L': L_,
            'W': W_,
            'H': H_,
            'm': mass,
            'v': v_,  # 0/1
            'x': x_,
            'y': y_,
            'z': z_
        }
        devices_info.append(dev_dict)
    return devices_info


########################################################################
#      2) 定义目标函数(甲板周长S、质心偏移Q)及约束(重叠U+最小距离)        #
########################################################################

def compute_objectives_and_constraints(solution, devices_info, L, W, H, d_min):
    """
    对一个解solution(长度=4*21)计算:
      目标1: S (外包矩形周长)
      目标2: Q (质心与形心偏移)
      overlapU: 重叠量(同层)
      distPenalty: 最小距离惩罚(同层)
    输入:
      - solution: 数组 shape=(84,)
        [x1,y1,z1,v1, x2,y2,z2,v2, ..., x21,y21,z21,v21]
      - devices_info: 23个设备信息(最后两个 22,23 不动), 前21个可动
    """
    local_devinfo = copy.deepcopy(devices_info)
    movable_num = 21

    # 先将solution映射到前21个设备
    for i in range(movable_num):
        x_val = solution[4 * i + 0]
        y_val = solution[4 * i + 1]
        z_val = solution[4 * i + 2]
        v_val = solution[4 * i + 3]

        local_devinfo[i]['x'] = x_val
        local_devinfo[i]['y'] = y_val
        local_devinfo[i]['z'] = z_val
        # orientation 二分(0或1)
        if v_val < 0.5:
            local_devinfo[i]['v'] = 0
        else:
            local_devinfo[i]['v'] = 1

    # ========== 计算外包矩形(甲板周长S) ==========
    xs_minus, xs_plus = [], []
    ys_minus, ys_plus = [], []

    total_mass = 0.0
    sum_mx, sum_my, sum_mz = 0.0, 0.0, 0.0

    for dev in local_devinfo:
        # 根据 v=0/1 判定“哪条是长，哪条是宽”
        if dev['v'] == 0:
            half_x = dev['L'] / 2.0
            half_y = dev['W'] / 2.0
        else:
            half_x = dev['W'] / 2.0
            half_y = dev['L'] / 2.0

        cx, cy = dev['x'], dev['y']
        xs_minus.append(cx - half_x)
        xs_plus.append(cx + half_x)
        ys_minus.append(cy - half_y)
        ys_plus.append(cy + half_y)

        # 计算质心
        total_mass += dev['m']
        sum_mx += dev['m'] * dev['x']
        sum_my += dev['m'] * dev['y']
        sum_mz += dev['m'] * dev['z']

    min_x, max_x = min(xs_minus), max(xs_plus)
    min_y, max_y = min(ys_minus), max(ys_plus)

    # 题意: S = |minX| + |maxX| + |minY| + |maxY|
    S = abs(min_x) + abs(max_x) + abs(min_y) + abs(max_y)

    # ========== 计算质心和形心偏移Q ==========
    x0 = sum_mx / total_mass
    y0 = sum_my / total_mass
    z0 = sum_mz / total_mass

    xc = (min_x + max_x) / 2.0
    yc = (min_y + max_y) / 2.0
    zc = H / 2.0  # 按题意

    Q = math.sqrt((x0 - xc) ** 2 + (y0 - yc) ** 2 + (z0 - zc) ** 2)

    # ========== 计算重叠量 overlapU ==========
    overlapU = 0.0
    n = len(local_devinfo)
    for i in range(n - 1):
        for j in range(i + 1, n):
            # 同层才检查
            if abs(local_devinfo[i]['z'] - local_devinfo[j]['z']) < 1e-6:
                # 取半长宽
                if local_devinfo[i]['v'] == 0:
                    p1_i = local_devinfo[i]['L'] / 2.0
                    p2_i = local_devinfo[i]['W'] / 2.0
                else:
                    p1_i = local_devinfo[i]['W'] / 2.0
                    p2_i = local_devinfo[i]['L'] / 2.0

                if local_devinfo[j]['v'] == 0:
                    p1_j = local_devinfo[j]['L'] / 2.0
                    p2_j = local_devinfo[j]['W'] / 2.0
                else:
                    p1_j = local_devinfo[j]['W'] / 2.0
                    p2_j = local_devinfo[j]['L'] / 2.0

                dx = (p1_i + p1_j) - abs(local_devinfo[i]['x'] - local_devinfo[j]['x'])
                dy = (p2_i + p2_j) - abs(local_devinfo[i]['y'] - local_devinfo[j]['y'])
                if dx > 0 and dy > 0:
                    overlapU += (dx + dy)
                    # print(local_devinfo[i],local_devinfo[j])

    # ========== 计算最小距离惩罚 distPenalty ==========
    distPenalty = 0.0
    for i in range(n - 1):
        for j in range(i + 1, n):
            if abs(local_devinfo[i]['z'] - local_devinfo[j]['z']) < 1e-6:
                if local_devinfo[i]['v'] == 0:
                    p1_i = local_devinfo[i]['L'] / 2.0
                    p2_i = local_devinfo[i]['W'] / 2.0
                else:
                    p1_i = local_devinfo[i]['W'] / 2.0
                    p2_i = local_devinfo[i]['L'] / 2.0

                if local_devinfo[j]['v'] == 0:
                    p1_j = local_devinfo[j]['L'] / 2.0
                    p2_j = local_devinfo[j]['W'] / 2.0
                else:
                    p1_j = local_devinfo[j]['W'] / 2.0
                    p2_j = local_devinfo[j]['L'] / 2.0

                # 需求: |x_i - x_j| - (p1_i + p1_j) >= d_min
                dx_ = abs(local_devinfo[i]['x'] - local_devinfo[j]['x']) - (p1_i + p1_j)
                dy_ = abs(local_devinfo[i]['y'] - local_devinfo[j]['y']) - (p2_i + p2_j)
                if dx_ < d_min and dy_ < d_min:
                    distPenalty += 1.0

    return S, Q, overlapU, distPenalty


########################################################################
#              3) 定义PSO中使用的粒子类 & 非支配判断函数等               #
########################################################################

def is_dominated(objA, objB):
    """
    判断objA=(f1A,f2A)是否被 objB=(f1B,f2B)支配:
    若 f1B <= f1A 且 f2B <= f2A，且至少有一个严格 <，则 B 支配 A
    """
    if (objB[0] <= objA[0]) and (objB[1] <= objA[1]) and ((objB[0] < objA[0]) or (objB[1] < objA[1])):
        return True
    else:
        return False


class Particle:
    def __init__(self, dim):
        self.position = np.zeros(dim)
        self.velocity = np.zeros(dim)
        self.pbest_position = np.zeros(dim)

        self.fitness = None  # (S,Q)
        self.fitness_penalty = None  # (overlapU, distPenalty)
        self.pbest_fitness = None
        self.pbest_penalty = None


########################################################################
#         4) 关键: 随机生成 x 时，使用并集区间 & 离散z处理              #
########################################################################

def random_x(a, b):
    """
    需求: x ∈ [-(a+b/2), -b/2) ∪ (b/2, b/2+a]
    算法:
      - 先随机抽 r ∈ [0,1)
      - 若 r < 0.5 => x ∈ [-(a+b/2), -b/2)
      - 否则 => x ∈ (b/2, b/2+a]
    返回在该并集中随机生成的 x
    """
    left_min = -(a + b / 2.0)
    left_max = -b / 2.0
    right_min = b / 2.0
    right_max = b / 2.0 + a

    r = random.random()
    if r < 0.5:
        return random.uniform(left_min, left_max)
    else:
        return random.uniform(right_min, right_max)


def repair_x(x_val, a, b):
    """
    若 x_val 落到禁区 (-b/2,b/2)，则将其挪到最近边界(或随机抉择)。
    你也可改成“重新生成”:
    """
    left_min = -(a + b / 2.0)
    left_max = -b / 2.0
    right_min = b / 2.0
    right_max = b / 2.0 + a

    if x_val >= left_max and x_val <= right_min:
        # 落在禁区
        # 简单做: 看哪个边界更近就贴哪
        dist_left = abs(x_val - left_max)
        dist_right = abs(x_val - right_min)
        if dist_left < dist_right:
            return left_max
        else:
            return right_min
    else:
        # 在 [left_min,left_max] ∪ [right_min,right_max] 范围内夹紧
        if x_val < left_min:
            return left_min
        if x_val > right_max:
            return right_max
        return x_val

########################################################################
#              5) 多目标PSO主流程: 初始化 -> 迭代 -> 存档更新           #
########################################################################

def MOPSO(devices_info, max_iter, pop_size, L, W, H, C, c1, c2, a, b, d_min):
    """
    多目标PSO, 前21个设备可动, dim=21*4=84.
    x ∈ 并集区间, y ∈ [-W/2,W/2], z ∈ {0, 1*H/(C*c1), 2*H/(C*c1),...}, v∈ {0或1}.
    """
    movable_num = 21
    dim = movable_num * 4

    # z离散步长
    z_step = H / float(C * c1)  # 例如 36/(2*3)=36/6=6
    possible_z_values = [k * z_step for k in range(C * c1 + 1)]  # [0,6,12,18,24,30,36?]

    # y区间
    min_y = -W / 2.0
    max_y = W / 2.0

    # 外部存档
    archive_positions = []
    archive_fitness = []
    archive_penalty = []

    # ============ 初始化粒子群 =============
    swarm = []
    for i in range(pop_size):
        p = Particle(dim)

        # 随机初始化 position
        for d in range(movable_num):
            # x
            p.position[4 * d + 0] = random_x(a, b)
            # y
            p.position[4 * d + 1] = random.uniform(min_y, max_y)
            # z
            p.position[4 * d + 2] = random.choice(possible_z_values)
            # v(orientation) 用 [0,1)表示，再二分
            p.position[4 * d + 3] = random.random()

        # 速度初始为0
        p.velocity = np.zeros(dim)

        # 计算适应度
        s_val, q_val, overlapU, distPen = compute_objectives_and_constraints(
            p.position, devices_info, L, W, H, d_min
        )
        p.fitness = (s_val, q_val)
        p.fitness_penalty = (overlapU, distPen)

        p.pbest_position = p.position.copy()
        p.pbest_fitness = (s_val, q_val)
        p.pbest_penalty = (overlapU, distPen)

        # 存档更新(如果无重叠,距离满足)
        if overlapU == 0 and distPen == 0:
            dominated_by_arch = False
            to_remove = []
            for idx_a in range(len(archive_positions)):
                if is_dominated(archive_fitness[idx_a], (s_val, q_val)):
                    to_remove.append(idx_a)
                elif is_dominated((s_val, q_val), archive_fitness[idx_a]):
                    dominated_by_arch = True
                    break

            for r in reversed(to_remove):
                archive_positions.pop(r)
                archive_fitness.pop(r)
                archive_penalty.pop(r)

            if not dominated_by_arch:
                archive_positions.append(p.position.copy())
                archive_fitness.append((s_val, q_val))
                archive_penalty.append((overlapU, distPen))

        swarm.append(p)

    # ============ 设置PSO参数 =============
    w = 0.5
    c1_ = 1.5
    c2_ = 1.5

    # ============ PSO 迭代 ============
    for iteration in range(max_iter):
        print(f"\n=== 迭代 {iteration + 1}/{max_iter}，当前存档规模={len(archive_positions)} ===")

        for i in range(pop_size):
            # 选全局引导 (从存档随机选一个)
            if len(archive_positions) > 0:
                g_idx = random.randint(0, len(archive_positions) - 1)
                gbest_pos = archive_positions[g_idx]
            else:
                gbest_pos = swarm[i].pbest_position

            # 更新速度 & 位置
            r1 = random.random()
            r2 = random.random()

            swarm[i].velocity = w * swarm[i].velocity \
                                + c1_ * r1 * (swarm[i].pbest_position - swarm[i].position) \
                                + c2_ * r2 * (gbest_pos - swarm[i].position)

            swarm[i].position = swarm[i].position + swarm[i].velocity

            # 越界/修复
            for d in range(movable_num):
                # x
                x_new = swarm[i].position[4 * d + 0]
                x_new = repair_x(x_new, a, b)  # 修复到并集区间
                swarm[i].position[4 * d + 0] = x_new

                # y
                y_new = swarm[i].position[4 * d + 1]
                if y_new < min_y:
                    y_new = min_y
                elif y_new > max_y:
                    y_new = max_y
                swarm[i].position[4 * d + 1] = y_new

                # z
                z_val = swarm[i].position[4 * d + 2]
                # 离散
                idx_z = round(z_val / (z_step + 1e-14))  # 避免除0
                idx_z = max(0, min(idx_z, len(possible_z_values) - 1))
                swarm[i].position[4 * d + 2] = possible_z_values[idx_z]

                # orientation
                v_val = swarm[i].position[4 * d + 3]
                if v_val < 0:
                    swarm[i].position[4 * d + 3] = 0
                elif v_val > 1:
                    swarm[i].position[4 * d + 3] = 1

            # 计算目标
            s_val, q_val, overlapU, distPen = compute_objectives_and_constraints(
                swarm[i].position, devices_info, L, W, H, d_min
            )
            swarm[i].fitness = (s_val, q_val)
            swarm[i].fitness_penalty = (overlapU, distPen)

            # print(f"粒子{i + 1}: S={s_val:.2f}, Q={q_val:.2f}, Overlap={overlapU:.2f}, Dist={distPen:.2f}")

            # 更新 pbest
            old_s, old_q = swarm[i].pbest_fitness
            old_ov, old_dp = swarm[i].pbest_penalty
            # 若新解无重叠/距离OK，且不被旧pbest支配，则更新
            if overlapU == 0 and distPen == 0:
                if not is_dominated((old_s, old_q), (s_val, q_val)):
                    swarm[i].pbest_position = swarm[i].position.copy()
                    swarm[i].pbest_fitness = (s_val, q_val)
                    swarm[i].pbest_penalty = (overlapU, distPen)

            # 更新存档
            if overlapU == 0 and distPen == 0:
                dominated_by_arch = False
                to_remove = []
                for idx_a in range(len(archive_positions)):
                    if is_dominated(archive_fitness[idx_a], (s_val, q_val)):
                        to_remove.append(idx_a)
                    elif is_dominated((s_val, q_val), archive_fitness[idx_a]):
                        dominated_by_arch = True
                        break
                for r in reversed(to_remove):
                    archive_positions.pop(r)
                    archive_fitness.pop(r)
                    archive_penalty.pop(r)

                if not dominated_by_arch:
                    archive_positions.append(swarm[i].position.copy())
                    archive_fitness.append((s_val, q_val))
                    archive_penalty.append((overlapU, distPen))

        # 每代结束可打印一下最优解
        if len(archive_positions) > 0:
            best_index = np.argmin([f[0] + f[1] for f in archive_fitness])  # S+Q最小
            print(
                f"  本代结束: 存档规模={len(archive_positions)}, S+Q最小解=({archive_fitness[best_index][0]:.2f},{archive_fitness[best_index][1]:.2f})")
        else:
            print("  本代结束: 存档依然为空(尚无可行解)")

    return archive_positions, archive_fitness, archive_penalty


########################################################################
#      6) 画图         #
########################################################################

def plot_nondominated_solutions(archive_fitness):
    """
    给定存档中的 (S,Q) 列表, 用matplotlib画散点
    """
    s_vals = [f[0] for f in archive_fitness]
    q_vals = [f[1] for f in archive_fitness]

    plt.figure(figsize=(6, 5))
    plt.scatter(s_vals, q_vals, c='r', marker='o', label='非劣解')
    plt.xlabel('S (甲板周长)')
    plt.ylabel('Q (质心偏移量)')
    plt.title('非劣解分布图')
    plt.legend()
    plt.grid(True)
    plt.show()


########################################################################
#                7) 主函数: 综合读取->PSO->结果输出->可视化              #
########################################################################

if __name__ == "__main__":
    # 1) 读取初始参数
    params_file = "./初始参数.txt"
    max_iter, pop_size, L, W, H, C, c1, c2, a, b, d_min = read_initial_params(params_file)
    print("读取到参数: max_iter={}, pop_size={}, L={}, W={}, H={}, C={}, c1={}, c2={}, a={}, b={}, d_min={}".format(
        max_iter, pop_size, L, W, H, C, c1, c2, a, b, d_min
    ))

    # 2) 读取设备数据(23个)
    data_file = "数据.txt"
    devices_info = read_devices_data(data_file)
    print(f"共读取到设备数量: {len(devices_info)} (其中后2个固定).")

    # 3) 先计算“原始txt坐标”那套解(若要看其目标函数、是否可行)
    #    把前21个当做可动，但是此处只做一次“测试计算”(看它 overlapU等)
    original_sol = np.zeros((21 * 4,))
    for i in range(21):
        original_sol[4 * i + 0] = devices_info[i]['x']
        original_sol[4 * i + 1] = devices_info[i]['y']
        original_sol[4 * i + 2] = devices_info[i]['z']
        # orientation保持 txt 中的(0 or 1), 先转成 0.0 或 1.0
        original_sol[4 * i + 3] = float(devices_info[i]['v'])
        # print(original_sol[4 * i + 0],original_sol[4 * i + 1],original_sol[4 * i + 2])
    s0, q0, ov0, dp0 = compute_objectives_and_constraints(original_sol, devices_info, L, W, H, d_min)
    print("【原始txt解】 S={:.3f}, Q={:.3f}, overlap={:.3f}, distPenalty={:.3f}".format(s0, q0, ov0, dp0))

    # 4) 执行MOPSO
    archive_positions, archive_fitness, archive_penalty = MOPSO(
        devices_info, max_iter, pop_size,
        L, W, H, C, c1, c2, a, b, d_min
    )

    # 5) 输出非劣解 + 可视化
    print("\n======== 优化完成，非劣解数量:{} ========".format(len(archive_positions)))
    with open("优化结果输出.txt", "w", encoding="utf-8") as fw:
        fw.write(f"非劣解数量: {len(archive_positions)}\n")
        for i, sol in enumerate(archive_positions):
            fS, fQ = archive_fitness[i]
            ov, dp = archive_penalty[i]
            fw.write(f"\n--- 解 {i + 1} ---\n")
            fw.write(f"S={fS:.3f}, Q={fQ:.3f}, Overlap={ov:.3f}, DistPenalty={dp:.3f}\n")

            # 写出21个可移动设备(其余2个固定)
            for j in range(21):
                xx = sol[4 * j + 0]
                yy = sol[4 * j + 1]
                zz = sol[4 * j + 2]
                vv = sol[4 * j + 3]
                fw.write(
                    f"设备ID={devices_info[j]['id']} x={xx:.3f}, y={yy:.3f}, z={zz:.3f}, v={0 if vv < 0.5 else 1}\n")

    # 若想画 (S,Q) 的分布图:
    if len(archive_positions) > 0:
        plot_nondominated_solutions(archive_fitness)

